
#ifndef AFTERMATH_PROBABILITY_EMPIRICAL_MEASURE_HPP_INCLUDED
#define AFTERMATH_PROBABILITY_EMPIRICAL_MEASURE_HPP_INCLUDED

#include "../type_concepts.hpp"
#include "dictionary.hpp"
#include "statistics_builder.hpp"

#include <cmath>
#include <cstdint>
#include <functional>
#include <iostream>
#include <string>
#include <system_error>
#include <type_traits>

namespace aftermath
{
    namespace probability
    {
        // TODO: hide <t_is_ordered> and <t_is_linear_space> parameters (possibly inside inheritance?) so that they are not accessible to the user.
        template <typename t_key_type, 
            typename t_count_type = std::size_t,
            typename t_sum_type = typename type_impl::get_product_result<t_key_type, t_count_type>::type,
            //typename t_sum_type = typename type_impl::get_sum_result<t_key_type, t_count_type>::type,
            typename t_mean_type = typename type_impl::get_divide_result<t_sum_type, double>::type,
            typename t_probability_type = double,
            bool t_is_ordered = detail::order_statistic<t_key_type, t_probability_type>::is_enabled,
            bool t_is_linear_space = detail::linear_statistic<t_key_type, t_count_type, t_sum_type, t_mean_type>::is_enabled,
            bool t_has_variance = detail::variance_statistic<t_key_type, t_count_type, t_sum_type, t_mean_type>::is_enabled>
        struct empirical_measure
        {
            static const bool is_ordered = t_is_ordered;
            static const bool is_linear_space = t_is_linear_space;
            static const bool has_variance = t_has_variance;

            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            typedef typename detail::dictionary<key_type, count_type>::type dictionary_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            template <typename, typename, typename, typename, typename, bool, bool, bool>
            friend struct empirical_measure;

        private:
            dictionary_type m_data = dictionary_type();
            detail::order_statistic<key_type, probability_type> m_order_statistic;
            detail::linear_statistic<key_type, count_type, sum_type, mean_type> m_linear_statistic;
            //detail::variance_statistic<key_type, count_type, sum_type, mean_type> m_variance_statistic;

            count_type m_count_observations = count_type();
            count_type m_max_height = count_type();
            key_type m_most_likely_value;

        public:
            empirical_measure()
                : m_order_statistic(), m_linear_statistic()//, m_variance_statistic()
            {
                // <t_key_type> has to implement std::hash and equality.
                static_assert(type_impl::has_equality<key_type>::value && type_impl::has_hash<key_type>::value,
                    "<t_key_type> has to implement binary '==' and std::hash.");
            }

            template <typename o_key_type, typename o_sum_type, typename o_mean_type, typename o_probability_type, bool o_is_ordered, bool o_is_linear_space, bool o_has_variance, typename t_transformer_type>
            empirical_measure(
                const empirical_measure<o_key_type, count_type, o_sum_type, o_mean_type, o_probability_type, o_is_ordered, o_is_linear_space, o_has_variance>& other,
                const t_transformer_type& transformer)
                : m_order_statistic(), m_linear_statistic()//, m_variance_statistic()
            {
                static_assert(type_impl::has_equality<key_type>::value && type_impl::has_hash<key_type>::value,
                    "<t_key_type> has to implement binary '==' and std::hash.");
                for (const auto& item : other.m_data) this->observe(transformer(item.first), item.second);
            }

            void clear()
            {
                this->m_data.clear();
                this->m_count_observations = count_type();
                this->m_max_height = count_type();

                this->m_order_statistic.clear();
                this->m_linear_statistic.clear();
                //this->m_variance_statistic.clear();
            }

            template <typename o_key_type, typename o_sum_type, typename o_mean_type, typename o_probability_type, bool o_is_ordered, bool o_is_linear_space, bool o_has_variance, typename t_transformer_type>
            void merge(
                const empirical_measure<o_key_type, count_type, o_sum_type, o_mean_type, o_probability_type, o_is_ordered, o_is_linear_space, o_has_variance>& other,
                const t_transformer_type& transformer)
            {
                detail::dictionary<key_type, count_type>::merge(this->m_data, other.m_data, transformer);
                this->rebuild_statistic();
            }

            template <typename o_sum_type, typename o_mean_type, typename o_probability_type, bool o_is_ordered, bool o_is_linear_space, bool o_has_variance>
            void merge(const empirical_measure<key_type, count_type, o_sum_type, o_mean_type, o_probability_type, o_is_ordered, o_is_linear_space, o_has_variance>& other)
            {
                detail::dictionary<key_type, count_type>::merge(this->m_data, other.m_data);
                this->rebuild_statistic();
            }

            count_type count_occurences(const key_type& key) const
            {
                return detail::dictionary<key_type, count_type>::find_or_default(this->m_data, key, count_type());
            }

            probability_type operator[](const key_type& key) const
            {
                return detail::dictionary<key_type, count_type>::find_or_default(this->m_data, key, count_type()) / static_cast<probability_type>(this->m_count_observations);
            }

            void copy_to(key_type* p_keys, count_type* p_counts) const
            {
                detail::dictionary<key_type, count_type>::copy(this->m_data, p_keys, p_counts);
            }

            void copy_to_normalized(key_type* p_keys, probability_type* p_probabilities) const
            {
                detail::dictionary<key_type, count_type>::copy(this->m_data, p_keys, p_probabilities, static_cast<probability_type>(this->m_count_observations));
            }

            void observe(const key_type& key, count_type repeat = 1)
            {
                if (repeat == 0) return;
                count_type new_height = repeat;

                auto it = this->m_data.find(key);
                bool is_contained = it != this->m_data.end();
                if (is_contained) // Already contains key.
                {
                    new_height += (it->second); // new_height = current_height + repeat.
                    it->second = new_height;
                }

                bool is_new_min = false;
                bool is_new_max = false;
                this->update_statistic(key, repeat, new_height, is_new_min, is_new_max);

                if (is_contained) return; // No need to insert a new key.
                
                if (is_new_min)
                    this->m_data.insert(this->m_data.begin(), std::pair<t_key_type, count_type>(key, new_height));
                else if (is_new_max)
                    this->m_data.insert(this->m_data.end(), std::pair<t_key_type, count_type>(key, new_height));
                else 
                    this->m_data.insert(std::pair<t_key_type, count_type>(key, new_height));
            }

            count_type count_observations() const
            {
                return this->m_count_observations;
            }

            bool empty() const
            {
                return this->m_count_observations == 0;
            }

            probability_type get_max_probability() const
            {
                return this->m_max_height / static_cast<probability_type>(this->m_count_observations);
            }

            count_type get_most_likely_count() const
            {
                return this->m_max_height;
            }

            const key_type& get_most_likely_value() const
            {
                return this->m_most_likely_value;
            }

            const dictionary_type& get_data() const
            {
                return this->m_data;
            }

            friend std::ostream& operator <<(std::ostream& os, const type& that)
            {
                //if (that.m_count_observations == 0U) return os << "{}";
                that.print_to(os);
                return os;
            }

        private:
            void update_statistic(const key_type& key, count_type repeat, count_type new_height, bool& is_new_min, bool& is_new_max)
            {
                this->m_order_statistic.observe(key, std::ref(is_new_min), std::ref(is_new_max));
                this->m_linear_statistic.observe(key, repeat);
                //this->m_variance_statistic.observe(key, repeat);

                if (this->m_max_height < new_height)
                {
                    this->m_max_height = new_height;
                    this->m_most_likely_value = key;
                }
                this->m_count_observations += repeat;
            }

            void rebuild_statistic()
            {
                this->m_order_statistic.clear();
                this->m_linear_statistic.clear();
                //this->m_variance_statistic.clear();

                bool is_new_min = false;
                bool is_new_max = false;
                this->m_count_observations = count_type();
                this->m_max_height = count_type();
                for (const auto& item : this->m_data) this->update_statistic(item.first, item.second, item.second, is_new_min, is_new_max);
            }

            template <typename t_stream_type>
            typename std::enable_if<std::is_integral<key_type>::value && type_impl::has_left_shift<t_stream_type, key_type>::value>::type
                print_to(t_stream_type& os, std::size_t min_height = 0UL, std::size_t max_height = 20UL) const
            {
                count_type count = this->m_count_observations;
                if (count == 0) os << "{}";
                else
                {
                    probability_type scale = this->get_max_probability();
                    for (key_type key = this->m_order_statistic.get_min(); key <= this->m_order_statistic.get_max(); ++key)
                    {
                        probability_type p = this->operator[](key);
                        std::size_t height = min_height + static_cast<std::size_t>((p / scale) * (max_height - min_height));
                        os << key << "\t" << std::string(height, '.') << std::string((1 + max_height) - height, ' ') << (100 * p) << '%' << std::endl;
                    }
                }
            }

            template <typename t_stream_type>
            typename std::enable_if<!std::is_integral<key_type>::value && type_impl::has_left_shift<t_stream_type, key_type>::value>::type
                print_to(t_stream_type& os) const
            {
                count_type count = this->m_count_observations;
                if (count == 0) os << "{}";
                else
                {
                    probability_type norm = static_cast<probability_type>(count);
                    for (const auto& item : this->m_data)
                    {
                        probability_type p = item.second / norm;
                        os << "{" << item.first << " : " << (100 * p) << "%}" << std::endl;
                    }
                }
            }

            template <typename t_stream_type>
            typename std::enable_if<!type_impl::has_left_shift<t_stream_type, key_type>::value>::type
                print_to(t_stream_type& os) const
            {
                count_type count = this->m_count_observations;
                if (count == 0) os << "{}";
                else os << "{...}";
            }
        };

        template <typename t_key_type, typename t_count_type, typename t_sum_type, typename t_mean_type, typename t_probability_type>
        struct empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, true, false, false>
            : public empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>
        {
            using empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>::empirical_measure;
            
            // Inherited typedef's:
            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            // Because of inheriting we should overwrite the two constants...
            static const bool is_ordered = true;
            static const bool is_linear_space = false;
            static const bool has_variance = false;
            // ...and a type definition.
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            // Now the extra functionality.
            const key_type& get_min() const { return this->m_order_statistic.get_min(); }
            const key_type& get_max() const { return this->m_order_statistic.get_max(); }
            probability_type cdf(const key_type& key) const { return this->m_order_statistic.cdf(key, this->m_count_observations, this->m_data); }
            const key_type& percentile(probability_type probability) const { return this->m_order_statistic.percentile(probability, this->m_count_observations, this->m_data); }
        };
        
        template <typename t_key_type, typename t_count_type, typename t_sum_type, typename t_mean_type, typename t_probability_type>
        struct empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, true, false>
            : public empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>
        {
            using empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>::empirical_measure;
            
            // Inherited typedef's:
            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            // Because of inheriting we should overwrite the two constants...
            static const bool is_ordered = false;
            static const bool is_linear_space = true;
            static const bool has_variance = false;
            // ...and a type definition.
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            // Now the extra functionality.
            const sum_type& get_sum() const { return this->m_linear_statistic.get_sum(); }
            mean_type get_mean() const { return this->m_linear_statistic.get_mean(this->m_count_observations); }
        };

        template <typename t_key_type, typename t_count_type, typename t_sum_type, typename t_mean_type, typename t_probability_type>
        struct empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, true, true>
            : public empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>
        {
            using empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>::empirical_measure;
            
            // Inherited typedef's:
            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            // Because of inheriting we should overwrite the two constants...
            static const bool is_ordered = false;
            static const bool is_linear_space = true;
            static const bool has_variance = true;
            // ...and a type definition.
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            // Now the extra functionality.
            const sum_type& get_sum() const { return this->m_linear_statistic.get_sum(); }
            mean_type get_mean() const { return this->m_linear_statistic.get_mean(this->m_count_observations); }
            // ...and some more.
            double compute_variance() const { return variance_helper_type::compute_variance(this->m_count_observations, this->m_data, this->get_mean()); }
            double compute_standard_deviation() const { return variance_helper_type::compute_standard_deviation(this->m_count_observations, this->m_data, this->get_mean()); }
        };

        template <typename t_key_type, typename t_count_type, typename t_sum_type, typename t_mean_type, typename t_probability_type>
        struct empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, true, true, false>
            : public empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>
        {
            using empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>::empirical_measure;
            
            // Inherited typedef's:
            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            // Because of inheriting we should overwrite the two constants...
            static const bool is_ordered = true;
            static const bool is_linear_space = true;
            static const bool has_variance = false;
            // ...and a type definition.
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            // Now the extra functionality...
            const key_type& get_min() const { return this->m_order_statistic.get_min(); }
            const key_type& get_max() const { return this->m_order_statistic.get_max(); }
            probability_type cdf(const key_type& key) const { return this->m_order_statistic.cdf(key, this->m_count_observations, this->m_data); }
            const key_type& percentile(probability_type probability) const { return this->m_order_statistic.percentile(probability, this->m_count_observations, this->m_data); }
            // ...and some more.
            const sum_type& get_sum() const { return this->m_linear_statistic.get_sum(); }
            mean_type get_mean() const { return this->m_linear_statistic.get_mean(this->m_count_observations); }
        };

        template <typename t_key_type, typename t_count_type, typename t_sum_type, typename t_mean_type, typename t_probability_type>
        struct empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, true, true, true>
            : public empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>
        {
            using empirical_measure<t_key_type, t_count_type, t_sum_type, t_mean_type, t_probability_type, false, false, false>::empirical_measure;
            
            // Inherited typedef's:
            typedef t_key_type key_type;
            typedef t_count_type count_type;
            typedef t_sum_type sum_type;
            typedef t_mean_type mean_type;
            typedef t_probability_type probability_type;
            typedef detail::variance_statistic<key_type, count_type, sum_type, mean_type> variance_helper_type;

            // Because of inheriting we should overwrite the two constants...
            static const bool is_ordered = true;
            static const bool is_linear_space = true;
            static const bool has_variance = true;
            // ...and a type definition.
            typedef empirical_measure<key_type, count_type, sum_type, mean_type, probability_type, is_ordered, is_linear_space, has_variance> type;

            // Now the extra functionality...
            const key_type& get_min() const { return this->m_order_statistic.get_min(); }
            const key_type& get_max() const { return this->m_order_statistic.get_max(); }
            probability_type cdf(const key_type& key) const { return this->m_order_statistic.cdf(key, this->m_count_observations, this->m_data); }
            const key_type& percentile(probability_type probability) const { return this->m_order_statistic.percentile(probability, this->m_count_observations, this->m_data); }
            // ...and some more...
            const sum_type& get_sum() const { return this->m_linear_statistic.get_sum(); }
            mean_type get_mean() const { return this->m_linear_statistic.get_mean(this->m_count_observations); }
            // ...and some more.
            double compute_variance() const { return variance_helper_type::compute_variance(this->m_count_observations, this->m_data, this->get_mean()); }
            double compute_standard_deviation() const { return variance_helper_type::compute_standard_deviation(this->m_count_observations, this->m_data, this->get_mean()); }
        };
    }
}

#endif // AFTERMATH_PROBABILITY_EMPIRICAL_MEASURE_HPP_INCLUDED
